package leaky_token

import (
	"ratelimit"
	"sync"
	"time"
)

type LeakyTokenLimit struct {
	rate       float64 // 漏水速率, 单位是秒
	capacity   float64 // 桶的容量
	water      float64 // 当前水量
	lastLeakMs int64   // 上次漏水时间戳（毫秒）
	lock       sync.Mutex
}

func New(rate float64, capacity float64) ratelimit.Limiter {
	limiter := &LeakyTokenLimit{
		rate:       rate,
		capacity:   capacity,
		lastLeakMs: time.Now().UnixNano() / 1e6,
		lock:       sync.Mutex{},
	}

	return limiter
}

func (l *LeakyTokenLimit) Allow() bool {
	return l.tryWater(1)
}

func (l *LeakyTokenLimit) tryWater(water float64) bool {
	l.lock.Lock()
	defer l.lock.Unlock()

	now := time.Now().UnixNano() / 1e6
	millisecondsSinceLastLeak := now - l.lastLeakMs
	waterLeaked := float64(millisecondsSinceLastLeak) * l.rate / 1000.0

	l.water = max(0, l.water-waterLeaked)
	l.lastLeakMs = now

	if water+l.water <= l.capacity {
		l.water += water
		return true
	}
	return false
}

func max(a, b float64) float64 {
	if a > b {
		return a
	}
	return b
}
